from flask import Flask, request, jsonify
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import GenerationConfig
from qwen_generation_utils import make_context, decode_tokens
import time

app = Flask(__name__)

tokenizer = AutoTokenizer.from_pretrained('Qwen-Chat/Qwen-14B-Chat-int4',  padding_side='left', trust_remote_code=True) # pad_token='<|extra_0|>', eos_token='<|endoftext|>',
model = AutoModelForCausalLM.from_pretrained('Qwen-Chat/Qwen-14B-Chat-int4', pad_token_id=tokenizer.pad_token_id, device_map="auto", trust_remote_code=True).eval()
model.generation_config = GenerationConfig.from_pretrained('Qwen-Chat/Qwen-14B-Chat-int4', pad_token_id=tokenizer.pad_token_id)

response, history = model.chat(tokenizer, "你好", history=None)
print(response)


def generate_responses(question,history):
    _chatbot = []
    # response, _ = model.chat(tokenizer, question, history=history)
    for response in model.chat_stream(tokenizer, question, history=history):
        yield response

@app.route('/generate', methods=['POST'])
def generate():
    data = request.json
    question = data.get('question', '')
    history = data.get('history', [])
    responses = generate_responses(question,history)
    for response in responses:
        print(response)
        yield jsonify(response)

if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0', port=6690)
